# knowledge-library/main.py
import os
import logging
from fastapi import FastAPI, UploadFile, File, HTTPException, BackgroundTasks
from fastapi.responses import JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from pathlib import Path
from typing import List
from config import settings
from services import (
    DocumentProcessor,
    VectorDBManager,
    OpenRouterClient,
    QueryProcessor
)
from services.models import (
    UploadResponse,
    QueryRequest,
    QueryResponse,
    ChunkUpdateRequest
)

# 初始化应用
app = FastAPI(
    title="个人知识库系统",
    version="1.0",
    docs_url="/api/docs",
    redoc_url="/api/redoc"
)

# CORS配置
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

# 初始化核心服务
processor = DocumentProcessor(
    chunk_size=settings.chunk_size,
    chunk_overlap=settings.chunk_overlap
)
vector_db = VectorDBManager()
openrouter = OpenRouterClient()
# 正确初始化方式：
query_processor = QueryProcessor(
    vector_db=vector_db,
    llm_client=openrouter
)

# 日志配置
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


@app.post("/api/upload", response_model=UploadResponse)
async def upload_files(
        files: List[UploadFile] = File(...),
        background_tasks: BackgroundTasks = BackgroundTasks()
):
    """文件上传端点"""
    saved_paths = []
    errors = []
    upload_dir = Path("temp_uploads")
    upload_dir.mkdir(exist_ok=True)

    for file in files:
        try:
            # 验证文件大小
            file_size = 0
            file_path = upload_dir / file.filename

            # 流式写入文件并计算大小
            with open(file_path, "wb") as buffer:
                while content := await file.read(1024 * 1024):  # 1MB chunks
                    file_size += len(content)
                    if file_size > settings.max_file_size * 1024 * 1024:
                        raise HTTPException(
                            status_code=413,
                            detail=f"文件 {file.filename} 超过大小限制"
                        )
                    buffer.write(content)

            saved_paths.append(str(file_path))
        except Exception as e:
            errors.append(file.filename)
            logger.error(f"文件 {file.filename} 上传失败: {str(e)}")
            continue

    # 启动后台处理任务
    if saved_paths:
        background_tasks.add_task(process_documents, saved_paths)

    return UploadResponse(
        status="completed" if not errors else "partial",
        processed_files=[f.filename for f in files if f.filename not in errors],
        error_files=errors
    )


async def process_documents(file_paths: List[str]):
    """后台文档处理任务"""
    for path in file_paths:
        try:
            logger.info(f"开始处理文档：{path}")
            chunks = processor.process_file(path)
            vector_db.upsert_chunks(chunks)
            logger.info(f"成功处理 {len(chunks)} 个文档块")

            # 清理临时文件
            os.remove(path)
        except Exception as e:
            logger.error(f"文档处理失败 {path}: {str(e)}")


@app.post("/api/query", response_model=QueryResponse)
async def handle_query(request: QueryRequest):
    """混合查询端点"""
    try:
        results = await query_processor.hybrid_query(
            question=request.question,
            use_llm=request.use_llm,
            top_k=request.top_k
        )
        return QueryResponse(**results)
    except Exception as e:
        logger.error(f"查询处理失败：{str(e)}")
        raise HTTPException(
            status_code=500,
            detail="查询处理失败，请检查日志"
        )


@app.put("/api/chunks/{chunk_id}")
async def update_chunk(chunk_id: str, request: ChunkUpdateRequest):
    """更新文档块"""
    try:
        vector_db.update_chunk(
            chunk_id=chunk_id,
            new_content=request.content,
            metadata=request.metadata
        )
        return {"status": "success"}
    except Exception as e:
        logger.error(f"更新文档块失败：{str(e)}")
        raise HTTPException(
            status_code=500,
            detail="文档块更新失败"
        )


@app.get("/api/health")
async def health_check():
    """健康检查端点"""
    return {
        "status": "healthy",
        "services": {
            "chroma": vector_db.client.heartbeat()
        }
    }


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(
        app,
        host="0.0.0.0",
        port=8000,
        log_level="info"
    )